import os
from PIL import Image
import numpy as np

# 定义文件夹路径
folder_A = r'E:\CODE\unet_pytorch_lc\img_out\long_best_IQI'
folder_B = r'E:\CODE\unet_pytorch_lc\img_out\YOLOv5识别'

# 创建一个输出文件夹
output_folder = r'E:\CODE\unet_pytorch_lc\img_out\融合'
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 定义红色范围
lower_red = np.array([200, 0, 0])
upper_red = np.array([256, 0, 0])


def is_red(pixel):
    """判断一个像素是否在指定的红色范围内"""
    return lower_red[0] < pixel[0] < upper_red[0]


# 获取文件夹A中的所有图片文件
files = [f for f in os.listdir(folder_A) if f.endswith('.jpg') or f.endswith('.png')]
total_files = len(files)  # 总文件数

# 遍历文件夹A中的图片
for index, filename in enumerate(files):
    # 打开文件夹A和B中的图片
    img_A = Image.open(os.path.join(folder_A, filename))
    img_B = Image.open(os.path.join(folder_B, filename))

    # 将图片转换为numpy数组
    img_A_np = np.array(img_A)
    img_B_np = np.array(img_B)

    # 创建一个掩码，只保留红色部分
    mask = np.zeros_like(img_A_np[:, :, 0], dtype=bool)  # 只需要一个通道的掩码
    for i in range(img_A_np.shape[0]):
        for j in range(img_A_np.shape[1]):
            if is_red(img_A_np[i, j]):
                mask[i, j] = True

    # 用掩码将红色部分覆盖到图片B
    img_B_np[mask] = img_A_np[mask]

    # 将结果转换为图片并保存
    result_img = Image.fromarray(img_B_np)
    result_img.save(os.path.join(output_folder, filename))

    # 打印进度
    print(f"处理进度: {index + 1}/{total_files}")

print("图像处理完成，结果保存在输出文件夹中。")
